[https://nvbugs/6482566][fix] Make MoE all-to-all completion-flag timeout phase-aware - #17376

Description

The MoE all-to-all dispatch/combine kernels wait for every peer's completion flag under a hardcoded ~300 s device-side deadline and execute trap; when it expires, which poisons the CUDA context of every waiting rank.

No collective separates a rank's first-touch JIT/autotune work from its dispatch launch, so that device-side budget is in effect a deadline on the slowest peer's host-side progress through its first forward. During startup warmup a rank can legitimately spend minutes in one-time work (CuTe DSL compilation, autotuning, module loading). When one does, its peers trap — a survivable slow start becomes a fatal crash for the whole worker.

Reported symptom: N-1 ranks print
dispatch: ---Rank X timed out waiting for completion flag from rank Y
and fail with CUDA_ERROR_LAUNCH_FAILED, while the blamed rank shows no error at all.

Fixes nvbugs/6482566.

Changes

• Runtime completion-flag budget. The timeout becomes a value carried in the kernel params rather than a compile-time constant, resolved by moeA2AGetTimeoutCycles(). Steady state is unchanged at 300 s; warmup uses 1800 s, the same order of magnitude as torch.distributed's default collective timeout. Overridable via TRTLLM_MOE_A2A_TIMEOUT_SEC / TRTLLM_MOE_A2A_WARMUP_TIMEOUT_SEC.
• Phase selection. New trtllm::moe_a2a_set_warmup op, driven from the is_warmup setter — the single choke point every warmup transition passes through — so the elevated budget cannot latch on into steady state. Builds without the op degrade to the previous behaviour.
• CUDA-graph capture is wrapped so captured graphs bake in the steady-state budget. The timeout is a launch argument, so it would otherwise be frozen at warmup's relaxed value for the life of the process.
• Indexer-Q prewarm. The DSv4 indexer-Q CuTe DSL kernels are compiled before the first collective-bearing forward. This removes the largest single contributor to the skew but is a partial mitigation only — other first-touch compiles sit on the all-to-all path itself and cannot be pre-compiled this way, which is why the runtime budget is the primary fix.

Test Coverage

tests/unittest/_torch/misc/test_moe_a2a_warmup_timeout.py (new, CPU-only):

• the phase selector forwards the value to the op;
• a build without the op does not break startup;
• CUDA-graph capture selects steady state and restores warmup afterwards;
• regression: the is_warmup setter switches the budget in both directions, so it cannot latch on after warmup.

End-to-end validation used fault injection on GB200 and GB300: delaying one rank by 600 s inside the first warmup forward reproduces the reported signature exactly (N-1 peers timing out in dispatch, all naming the delayed rank) under the 300 s budget, and starts cleanly under the new warmup budget.

Notes for reviewers

• Steady-state behaviour is intentionally unchanged; only the warmup window and the new env overrides alter the deadline.
• The budget is passed as a kernel launch argument. A device-resident value would let the spin loop observe phase changes live and would be phase-accurate under CUDA graphs; that is a larger change to the kernel inner loop and is deliberately left as follow-up, with capture pinned to steady state in the meantime.
• Replacing trap; with a recoverable error path is a separate, larger piece of work: after a timeout the receive buffers are incomplete and the flag generation would need invalidating across ranks. Not attempted here.

PR Checklist

• PR title follows [JIRA/NVBugs ID][type] Summary
• Commit is DCO signed off
• New unit tests added for the fixed behaviour
• CI green (draft — pending)

Dev Engineer Review

• Added configurable MoE all-to-all timeout budgets.
• Uses 300 seconds in steady state and 1,800 seconds during warmup.
• Supports TRTLLM_MOE_A2A_TIMEOUT_SEC and TRTLLM_MOE_A2A_WARMUP_TIMEOUT_SEC.
• Added moe_a2a_set_warmup(bool) with compatibility handling for unavailable native bindings.
• Uses the steady-state timeout during CUDA-graph capture.
• Prewarms DeepSeek-V4 indexer-Q CuTe DSL kernels before collective-bearing forwards.
• Propagates timeout cycles through dispatch and combine parameters.
• Fault-injection validation covered delayed ranks on GB200 and GB300.
• Review should verify environment-value validation, default handling, atomic state transitions, and unavailable-binding behavior.

QA Engineer Review

• Added tests/unittest/_torch/misc/test_moe_a2a_warmup_timeout.py.
• Added:
  • test_set_warmup_forwards_value_to_op
  • test_missing_op_is_tolerated
  • test_capture_context_selects_steady_state_then_restores
  • test_is_warmup_setter_switches_budget_both_ways
• The test file has no matching entry in tests/integration/test_lists/, test-db/, or qa/.
• Verdict: needs follow-up.